package com.dl4jlpr;

import lombok.extern.slf4j.Slf4j;
import org.deeplearning4j.nn.conf.inputs.InputType;
import org.deeplearning4j.nn.conf.layers.samediff.SDLayerParams;
import org.deeplearning4j.nn.conf.layers.samediff.SameDiffLayer;
import org.deeplearning4j.nn.params.DefaultParamInitializer;
import org.deeplearning4j.nn.weights.WeightInit;
import org.nd4j.autodiff.samediff.SDVariable;
import org.nd4j.autodiff.samediff.SameDiff;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.api.ops.impl.layers.convolution.Pooling2D;
import org.nd4j.linalg.api.ops.impl.layers.convolution.config.Pooling2DConfig;

import java.util.Arrays;
import java.util.Map;
@Slf4j
public class MaxPool_Fake3d extends SameDiffLayer {

    private int[] dimensions1;
    private int[] dimensions2;
    private int[] kernelSize;
    private int[] strideSize;
    private long nIn;
    private long nOut;



    public MaxPool_Fake3d(MaxPool_Fake3d.Builder builder) {
        super(builder);

        this.nIn = builder.nIn;
        this.nOut = builder.nOut;
        this.dimensions1=builder.dimensions1.clone();
        this.dimensions2=builder.dimensions2.clone();
        this.strideSize=builder.strideSize.clone();
        this.kernelSize=builder.kernelSize.clone();
    }
    private MaxPool_Fake3d(){
        //No op (Jackson)
    }

    @Override
    public void defineParameters(SDLayerParams sdLayerParams) {
        sdLayerParams.addWeightParam(DefaultParamInitializer.WEIGHT_KEY, (int) nIn, 1);
    }

    @Override
    public void initializeParameters(Map<String, INDArray> map) {
        initWeights( (int) nIn, 1, weightInit,  map.get(DefaultParamInitializer.WEIGHT_KEY));
    }

    @Override
    public SDVariable defineLayer(SameDiff sameDiff, SDVariable layerInput, Map<String, SDVariable> paramTable, SDVariable mask) {
        //Float[] f=layerInput.castTo(DataType.FLOAT);
        log.info("MaxPool_Fake3d input shape:{}", Arrays.toString(layerInput.getShape()));
        Pooling2DConfig layer1PoolConfig = Pooling2DConfig.builder()
                .kH(kernelSize[1])  //filter height
                .kW(kernelSize[2])  //filter width
                .sH(strideSize[1])  //stride height
                .sW(strideSize[2])  //stride width
                .type(Pooling2D.Pooling2DType.MAX)//max-pooling
                .build();
        Pooling2DConfig layer2PoolConfig = Pooling2DConfig.builder()
                .kH(kernelSize[0])  //filter height
                .kW(kernelSize[0])  //filter width
                .sH(strideSize[1])  //stride height
                .sW(strideSize[0])  //stride width
                .type(Pooling2D.Pooling2DType.MAX)//max-pooling
                .build();
        SDVariable pool1 =sameDiff.cnn.maxPooling2d(layerInput, layer1PoolConfig);
        SDVariable t1=pool1.permute(dimensions1);
        SDVariable pool2 =sameDiff.cnn.maxPooling2d(t1, layer2PoolConfig);
        SDVariable t2=pool2.permute(dimensions2);
        //t.reshape(-1,64,20,90);
        return  t2;
    }

    @Override
    public InputType getOutputType(int i, InputType inputType) {
        long[] shape=inputType.getShape();
        log.info("MaxPool_Fake3d input shape:{} and the lay index:{}",Arrays.toString(shape),String.valueOf(i));
        switch ((int) shape[0]){
            case 64:
                return InputType.convolutionalFlat(20,90,64);
            case 128:
                return InputType.convolutionalFlat(18,44,64);
            case 256:
                return InputType.convolutionalFlat(16,21,64);
        }
        return null;
    }

    public static class Builder extends SameDiffLayer.Builder<MaxPool_Fake3d.Builder> {
        private int[] dimensions1;
        private int[] dimensions2;
        private int[] kernelSize;
        private int[] strideSize;
        private int nIn;
        private int nOut;
        private WeightInit weightInit;
        @Override
        public Builder weightInit(WeightInit weightInit) {
            this.weightInit = weightInit;
            return this;
        }

        public Builder nIn(int nIn){
            this.nIn = nIn;
            return this;
        }

        public Builder nOut(int nOut){
            this.nOut = nOut;
            return this;
        }
        public Builder dimensions1(int ... dimensions1){
            this.dimensions1=dimensions1.clone();
            return this;
        }
        public Builder dimensions2(int ... dimensions2){
            this.dimensions2=dimensions2.clone();
            return this;
        }
        public Builder kernelSize(int ... kernelSize){
            this.kernelSize=kernelSize.clone();
            return this;
        }
        public Builder strideSize(int ... strideSize){
            this.strideSize=strideSize.clone();
            return this;
        }
        @Override
        public MaxPool_Fake3d build() {
            return new MaxPool_Fake3d(this);
        }
    }

}
